/*
MIT License
Copyright (c) 2025 Chuck Fu <changlin.fu@gmail.com>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/
#pragma once
#ifndef CF_CRYPTO_SM2_H
#define CF_CRYPTO_SM2_H

#include "stdio.h"
#include "stdlib.h"
#include "stdint.h"
#include "stdbool.h"

// Turn On Yield CPU implementation
// #define CF_CRYPTO_XMEN_INTEGRATION

// =============================================================================
// SM2 Utility Struct - Bit String
// The bitstring struct stores bits in byte array, and keep track of the active
// bit counts.
// =============================================================================
typedef struct bitstring_st {
    uint32_t bit_len;
    uint32_t byte_len;
    uint32_t data_cap;
    uint8_t *data;
} sm2_bitstring;

// create bitstring (Wrap existing storage)
sm2_bitstring sm2_bitstring_wrap(
    uint8_t *buf,
    uint32_t buf_size,
    uint32_t data_bits
);
// copy values to bitstring
void sm2_bitstring_set_from_cstr(
    sm2_bitstring *result,
    const char *cstr
);
void sm2_bitstring_set_from_hex_str(
    sm2_bitstring *result,
    const char *hex_str
);
void sm2_bitstring_set_from_bytes(
    sm2_bitstring *result,
    const uint8_t *bytes,
    uint32_t byte_len
);
// get and set bit value
int sm2_bitstring_get_bit_value(
    const sm2_bitstring *bitstring,
    uint32_t bit_idx
);
void sm2_bitstring_set_bit_value(
    sm2_bitstring *result,
    uint32_t bit_idx,
    bool bit_val
);

// =============================================================================
// SM2 Utility Struct - BigInteger
// A self-contained implementation of 256 bit bigint with the operation allowed
// in the prime Field Fp.
// =============================================================================
typedef enum sm2_cmp_order_enum {
    SM2_CMP_ORDER_LESS    = -1,
    SM2_CMP_ORDER_EQUAL   = 0,
    SM2_CMP_ORDER_GREATER = 1,
} sm2_cmp_order;

// Defines the bigint type used for SM2
typedef struct uint256_st {
    uint32_t data[8];
} sm2_bigint;

// set bigint values
void sm2_bigint_set_from_bytes(sm2_bigint *res, const uint8_t *bytes);
void sm2_bigint_set_from_hex_str(sm2_bigint *res, const char *hex_str);
void sm2_bigint_set_from_u32(sm2_bigint *res, uint32_t u32val);
void sm2_bigint_set_from_other(sm2_bigint *res, const sm2_bigint *other);
void sm2_bigint_set_from_bitstring(sm2_bigint *res, const sm2_bitstring *bs);

// Convert the bigint to byte array. (32 bytes)
void sm2_bigint_to_bytes(sm2_bigint *x, uint8_t *bytes);

// compare a and b, the result is -1 (less), 0(same), 1(greater)
sm2_cmp_order sm2_bigint_cmp(const sm2_bigint *a, const sm2_bigint *b);

// check if its zero
bool sm2_bigint_is_zero(const sm2_bigint *x);


// -----------------------------------------------------------------------------
// Bigint Math Functions in Fp:
// It supports: Add, Sub, Mul, Inv, Pow operations.
// Note: Div can be implemented with the inverse mod, e.g. a / b = a * Inv(b)
// -----------------------------------------------------------------------------
// t = t % p,
// this normalized t to the range [0, p)
void sm2_bigint_norm(
    sm2_bigint *t,
    const sm2_bigint *p
);

// t = (a + b) % p
void sm2_bigint_add_mod_p(
    sm2_bigint *t,
    const sm2_bigint *a,
    const sm2_bigint *b,
    const sm2_bigint *p
);
// t = (a - b) % p
void sm2_bigint_sub_mod_p(
        sm2_bigint *t,
        const sm2_bigint *a,
        const sm2_bigint *b,
        const sm2_bigint *p
);
// t = (a * b) % p
void sm2_bigint_mul_mod_p(
        sm2_bigint *t,
        const sm2_bigint *a,
        const sm2_bigint *b,
        const sm2_bigint *p
);
// t = (a^b) % p
void sm2_bigint_pow_mod_p(
        sm2_bigint *t,
        const sm2_bigint *a,
        const sm2_bigint *b,
        const sm2_bigint *p
);
// t = a^-1 % p
void sm2_bigint_inv_mod_p(
        sm2_bigint *t,
        const sm2_bigint *a,
        const sm2_bigint *p
);

// =============================================================================
// SM2 Library APIs
// =============================================================================
/**
 * struct sm2_ctx_st SM2 Running Context
 *
 * SM2 use the ECC (Elliptic Curve Crypto) curve in the 256 bit prime field
 * with the following function: y^2 = x^3 + ax^2 + b
 *
 * @p: prime number, represent the field Fp
 * @a: the parameter a in the ECC equation
 * @b: the parameter b in the ECC equation
 * @n: the order of base point G
 * @Gx: x cord of base point G
 * @Gy: y cord of base point G
 * @Px: x cord of public key point P
 * @Py: x cord of public key point P
 * @d: private key used
 * @k: random number used in signature generation & data encryption
 */
typedef struct sm2_ctx_st {
    sm2_bigint p;  // field p
    sm2_bigint a;  // ECC parameter a
    sm2_bigint b;  // ECC parameter b
    sm2_bigint n;  // Order of base point G(x,y)
    sm2_bigint Gx; // x cord of base point G
    sm2_bigint Gy; // y cord of base point G
    sm2_bigint Px; // x cord of public key point P
    sm2_bigint Py; // y cord of public key point P
    sm2_bigint d;  // Private key
    sm2_bigint k;  // Random number k used in signature generation & data encryption
    void (* sm2_change_k)(struct sm2_ctx_st *ctx); // Function to generate K
} sm2_ctx;

/**
 * struct sm_signature_st
 * @r: the r value in the signature
 * @s: the s value in the signature
 */
typedef struct sm2_signature_st {
    uint8_t r[32];
    uint8_t s[32];
} sm2_signature;

/**
 * Yield CPU for long running loop, used in the Embedded System Task.
 */
void sm2_yield_cpu();

/**
 * Encode SM2 Signature to DER (Distinguished Encoding Rules)
 * The ASN.1 def of SM2 signature is
 *
 * SM2Signature :: = SEQUENCE {
 *     r INTEGER,
 *     s INTEGER
 * }
 *
 * @param signature the SM2 signature <r, s> pair
 * @return The encoded DER types
 */
void sm2_signature_to_der(uint8_t *der, const sm2_signature *signature);

/**
 * Get the shared SM2 Context, used in the Embedded System setting.
 */
sm2_ctx *sm2_ctx_shared();

// SM2 functions
/**
 * Init the SM2 context with default parameters values.
 *
 * @param ctx the ctx object used
 */
void sm2_ctx_init(sm2_ctx *ctx);

/**
 * Set the private key used in the context.
 *
 * @param ctx the ctx object used
 * @param private_key the 32 bytes private_key to set.
 */
void sm2_ctx_set_private_key(sm2_ctx *ctx, const char *private_key);

/**
 * Set the public key used in the context.
 *
 * @param ctx the ctx object used
 * @param public_key the 64 bytes public key is in the format PxPy
 */
void sm2_ctx_set_public_key(sm2_ctx *ctx, const char *public_key);

// SM2 ECC Solver
/**
 * struct sm_ecc_point - Represent a vector/point in ECC space.
 * Note: When used in Affine Space, z cord represent the infinity point when
 *       z is not zero.
 * The implementation detail is hidden from user, only use
 * sm2_ecc_point_new_from_xy(x, y) to create a new ECC point in user space.
 *
 * @x the x cord of point
 * @y the y cord of point
 * @z the z cord of point
 */
typedef struct sm2_ecc_point_st {
    sm2_bigint x, y, z;
} sm2_ecc_point;

/**
 * Create a 0 point in ECC space (the infinity point)
 *
 * @return the created point in ECC space
 */
void sm2_ecc_point_set_zero(sm2_ecc_point *p);

/**
 * Create a point in ECC space use (x,y).
 *
 * @param x the x cord of the point
 * @param y the y cord of the point
 * @return the created point in ECC space
 */
void sm2_ecc_point_set_xy(
    sm2_ecc_point *res,
    const sm2_bigint *x,
    const sm2_bigint *y
);

/**
 * Point ADD operation in ECC space.
 *
 * @param ctx the ctx object used
 * @param res the result of the add operation, res = P1 + P2
 * @param p1 the point used in the add operation
 * @param p2 the point used in the add operation
 */
void sm2_ecc_point_add(
    const sm2_ctx *ctx,
    sm2_ecc_point *res,
    const sm2_ecc_point *p1,
    const sm2_ecc_point *p2
);

/**
 * Point MUL k operation in ECC space.
 *
 * @param ctx the ctx object used
 * @param p the point used in the add operation
 * @param k the point used in the add operation
 *
 * @return the result of the sum of k p's in ECC space. res = kP
 */
void sm2_ecc_point_mul_k(
    const sm2_ctx *ctx,
    sm2_ecc_point *res,
    const sm2_ecc_point *p,
    const sm2_bigint *k
);

/**
 * Normalized Point to <x, y, 1>
 *
 * @return the created point in ECC space
 */
void sm2_ecc_point_norm(const sm2_ctx *ctx, sm2_ecc_point *p);

/**
 * Create the user distinguishing identifier string
 * Zstr = ENTLA ∥ IDA ∥ a ∥ b ∥ Gx ∥ Gy ∥ Px ∥ Py
 *
 * @param ctx the ctx object used
 * @param id_a the user provided string
 * @return the user distinguishing identifier string
 */
void sm2_build_zstr(const sm2_ctx *ctx, sm2_bitstring *res, const sm2_bitstring *id_a);

/**
 * Create the digest of za and message.
 *
 * @param za the user distinguishing identifier
 * @param message the message to sign
 * @return za = H256(za||message)
 */
void sm2_build_digest(
    sm2_bitstring *digest,
	const sm2_bitstring* za,
    const sm2_bitstring* message
);

/**
 * Random bigint integer k generation, used in sign process and data encryption.
 *
 * @param ctx the ctx object used
 */
void sm2_change_k(sm2_ctx *ctx);


/**
 * Generate an SM2 signature against the given message digest.
 *
 * @param ctx the ctx object used
 * @param digest the result of H256(Za||Message)
 * @return 0 means signature generate succeed, <0 means error happened.
 */
int sm2_sign_generate(
    const sm2_ctx *ctx,
    const sm2_bitstring *digest,
    sm2_signature *signature
);

/**
 * Verifies an SM2 signature against the given message digest.
 *
 * @param ctx the ctx object used
 * @param digest the result of H256(Za||Message)
 * @param signature the signature to verify
 * @return 0 means signature verify succeed, <0 means error happened.
 */
int sm2_sign_verify(
    const sm2_ctx *ctx,
    const sm2_bitstring *digest,
    const sm2_signature *signature
);

/**
 * Encrypt Message.
 *   C = C1 || C3 || C2
 *   |C1| = 1 + 32 + 32 = 65 bytes = 520 bits
 *   |C3| = 256 bits = 32 bytes
 *   |C2| = message bit lens = klen
 *
 *   The data cap for encrypted_message needs >= (776 + klen + 7) / 8 bytes
 *
 * @param ctx the ctx object used
 * @param plain_message the plain message
 * @param encrypted_message (out) the encrypted message
 * @return 0 means encryption succeed, <0 means error happened.
 */
int sm2_data_encrypt(
    const sm2_ctx *ctx,
    const sm2_bitstring *plain_message,
    sm2_bitstring *encrypted_message
);

/**
 * Decrypt Message.
 *
 * @param ctx the ctx object used
 * @param encrypted_message the encrypted message, in C1C3C2 format
 * @param plain_message (out) the plain message
 * @return 0 means decryption succeed, < 0 means error happened.
 */
int sm2_data_decrypt(
    const sm2_ctx *ctx,
    const sm2_bitstring *encrypted_message,
    sm2_bitstring *plain_message
);

// =============================================================================
// SM2 Debug interfaces (with sm2_show & sm2_write prefix)
// =============================================================================
// Print to std output
void sm2_show_bytes(const uint8_t *bytes, uint32_t byte_len, bool with_space);
void sm2_show_bitstring(const sm2_bitstring *bs);
void sm2_show_bigint(const sm2_bigint *b);
void sm2_show_ctx(const sm2_ctx *ctx);
void sm2_show_ecc_point(const char *point_name, const sm2_ecc_point *p);
void sm2_show_signature(const sm2_signature *signature);
// Write to files
void sm2_write_bytes(FILE *file, const uint8_t *bytes, uint32_t byte_len);
void sm2_write_bitstring(FILE *file, const sm2_bitstring *bs);
void sm2_write_bigint(FILE *file, const sm2_bigint *x);

#endif //CF_CRYPTO_SM2_H
